import PIL.Image as Image
import os
from datetime import datetime
import cv2
import numpy as np

# 打印当前时间到毫秒，用于判断代码的性能。如果性能太低了，那么就逐帧运行
print(datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3])  # 去掉最后三位毫秒后的零

def process_image(input_path, output_path=None):
    """
    处理图片，将RGB像素值低于100的重置为0，然后保存新图片

    参数:
        input_path: 输入图片的路径
        output_path: 输出图片的路径，默认为在输入文件名后加"_processed"
    """
    try:
        # 打开图片
        with Image.open(input_path) as img:
            # 确保图片是RGB模式
            if img.mode != 'RGB':
                img = img.convert('RGB')

            # 获取图片尺寸
            width, height = img.size

            # 加载像素数据
            pixels = img.load()

            # 遍历每个像素
            for i in range(width):
                for j in range(height):
                    # 获取当前像素的RGB值
                    r, g, b = pixels[i, j]

                    # 将低于100的值重置为0
                    # new_r = r if r >= 200 else 0
                    # new_g = g if g >= 200 else 0
                    if g <= 200:
                        new_g = 0
                    else:
                        new_g = g
                    # new_b = b if b >= 100 else 0

                    if r > 200 and g > 200:
                        if b <= 100:
                            new_b = 0
                        else:
                            new_b = b
                    else:
                        new_b = 0

                    if r <= 200:
                        new_r = 0
                    else:
                        new_r = r
                        # 官方插件，做了红色的圆圈流动效果，fix一下
                        new_g = 210
                        new_b = 140

                    # new_r = 0
                    # new_g = 0
                    # new_b = 0

                    # 设置新的像素值
                    pixels[i, j] = (new_r, new_g, new_b)

            # 如果未指定输出路径，自动生成
            if output_path is None:
                # 分离文件名和扩展名
                base, ext = os.path.splitext(input_path)
                output_path = f"{base}_processed{ext}"

            # 保存处理后的图片（only for debug）
            img.save(output_path)
            print(f"图片处理完成，已保存至: {output_path}")
            return output_path, img

    except Exception as e:
        print(f"处理图片时出错: {str(e)}")
        return None, None

def hasZFXorNot(output_path=None,outputBinImg=None):
    # 读取图像（替换为你的图片路径，调试，用路径，实际运行不会采用磁盘文件交互一次）
    # image = cv2.imread(output_path)
    # 读取图像（替换为你的图片路径，调试，用路径，实际运行不会采用磁盘文件交互一次）
    # 1. 先转换为NumPy数组
    numpy_image = np.array(outputBinImg)
    # 2. 转换颜色空间（PIL默认是RGB，OpenCV需要BGR）
    image = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
    if image is None:
        print("无法读取图像，请检查路径是否正确")
        exit()

    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化处理（调整阈值，突出光亮区域，这里根据实际情况调整，示例用150作为阈值）
    _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)

    cv2.imwrite('yhb_debug_abc1.jpg', binary)

    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 遍历轮廓，判断是否有近似正方形的轮廓
    has_square = False
    for contour in contours:
        # 计算轮廓的周长
        perimeter = cv2.arcLength(contour, True)
        # 多边形近似
        approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)
        # 如果近似后是四边形
        if len(approx) == 4:
            # 计算轮廓的外接矩形
            x, y, w, h = cv2.boundingRect(approx)
            if w < 20 or h < 20:
                # 过滤掉识别得太小的正方形，技能键就是固定的大小
                continue
            # 自己画个框，验证识别出来的区域，坐标，长宽这些，是否符合要求
            test001 = cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
            # cv2.imshow('Original Grayscale', test001)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
            cv2.imwrite('yhb_debug_abc2.jpg', test001)
            # 计算宽高比，接近1说明是正方形（可根据实际情况调整容差，这里容差设为0.2）
            aspect_ratio = float(w) / h
            if 0.8 < aspect_ratio < 1.2:
                # 计算轮廓的面积和外接矩形的面积比，接近1说明轮廓接近矩形（可调整容差，这里容差设为0.2）
                area = cv2.contourArea(contour)
                rect_area = w * h
                if 0.8 < area / rect_area < 1.2:
                    has_square = True
                    break

    if has_square:
        #图片中存在光亮的类似正方形的方框
        return True
    else:
        #图片中不存在光亮的类似正方形的方框
        return False

if __name__ == "__main__":
    # 这个函数用于判断图片中，是否存在高亮的正方形块（wow中用到）
    input_image_path = "./yhb001.png"  # 替换为你的输入图片路径
    output_image_path = "output_image.jpg"  # 替换为你想要保存的路径，或设为None

    # 预处理图片，降低像素复杂的维度，返回处理过的中间图片
    path1, binImg = process_image(input_image_path, output_image_path)
    # 判断是否图片中，是否存在高亮的正方形块（wow中用到）
    boolResult = hasZFXorNot(output_image_path,binImg)
    if boolResult:
        print("图片中存在光亮的类似正方形的方框")
    else:
        print("图片中不存在光亮的类似正方形的方框")
    # 打印当前时间到毫秒
    print(datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3])  # 去掉最后三位毫秒后的零
